package org.nimi317.web_gis.model;

import jakarta.validation.constraints.NotNull;
import lombok.extern.slf4j.Slf4j;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.util.ModelSerializer;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.api.preprocessor.NormalizerStandardize;
import org.nd4j.linalg.factory.Nd4j;
import org.nimi317.web_gis.data.AreaCodeList;
import org.nimi317.web_gis.exception.E;
import org.nimi317.web_gis.exception.RException;
import org.springframework.beans.BeanUtils;
import org.springframework.context.ApplicationListener;
import org.springframework.context.event.ContextClosedEvent;
import org.springframework.stereotype.Component;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;
import java.util.Set;
import java.util.function.Function;

/**
 * @author thunderobot
 */
@Component
@Slf4j
public class AirQualityPredict implements ApplicationListener<ContextClosedEvent> {

    private final String basePath = "C:\\Users\\thunderobot\\Desktop\\毕业设计\\web_gis\\src\\main\\resources\\main\\";

    private final MultiLayerNetwork model;

    private final NormalizerStandardize standardize;

    private final Set<String> abnormalArray = AreaCodeList.abnormalArray;

    public AirQualityPredict() throws IOException {
        this.model = ModelSerializer.restoreMultiLayerNetwork(new File(basePath + "model.zip"));
        this.standardize = new NormalizerStandardize();
        this.standardize.fitLabel(true);
        this.standardize.load(
                new File(basePath + "mean.bin"),
                new File(basePath + "std.bin"),
                new File(basePath + "label_mean.bin"),
                new File(basePath + "label_std.bin")
        );
    }

    private List<Data> predict(List<Data> data) {
        //将data转化为二维数组
        int size = data.size();
        double[][] result = new double[size][7];
        for (int i = 0; i < size; i++) {
            result[i][0] = parseOrInsert(data, i, Data::getAqi);
            result[i][1] = parseOrInsert(data, i, Data::getPm25);
            result[i][2] = parseOrInsert(data, i, Data::getPm10);
            result[i][3] = parseOrInsert(data, i, Data::getCo);
            result[i][4] = parseOrInsert(data, i, Data::getNo2);
            result[i][5] = parseOrInsert(data, i, Data::getO3);
            result[i][6] = parseOrInsert(data, i, Data::getSo2);
        }
        //对二维数组进行异常值填充
        try (INDArray input = Nd4j.create(result).transpose().reshape(1, 7, size)) {
            //标准化
            standardize.transform(input);
            INDArray output = model.output(input);
            //输出结果
            standardize.revertLabels(output);

            INDArray labels = output.reshape(7, size).transpose();
            List<Data> dataList = new LinkedList<>(data);
            for (float[] floatMatrix : labels.toFloatMatrix()) {
                Data e = new Data();
                e.setAqi(String.format("%.2f", floatMatrix[0]));
                e.setPm25(String.format("%.2f", floatMatrix[1]));
                e.setPm10(String.format("%.2f", floatMatrix[2]));
                e.setCo(String.format("%.2f", floatMatrix[3]));
                e.setNo2(String.format("%.2f", floatMatrix[4]));
                e.setO3(String.format("%.2f", floatMatrix[5]));
                e.setSo2(String.format("%.2f", floatMatrix[6]));
                dataList.add(e);
            }
            return dataList;
        }
    }

    private double parseOrInsert(List<Data> data, int index, Function<Data, String> getter) {
        try {
            return Double.parseDouble(getter.apply(data.get(index)));
        } catch (NumberFormatException e) {
            return fill(data, index, getter);
        }
    }

    private double fill(List<Data> result, int i, Function<Data, String> method) {
        double up = search(result, i, method, true);
        double down = search(result, i, method, false);
        if (up == 0 && down == 0) {
            throw new RException(E.AirQualityPredictException);
        }
        return (up + down) / ((up == 0 || down == 0) ? 1 : 2);
    }

    private double search(List<Data> data, int i, Function<Data, String> method, boolean up) {
        String applied;
        while (true) {
            applied = method.apply(data.get(i));
            if (!abnormalArray.contains(applied)) {
                break;
            }
            if (up) {
                i--;
            } else {
                i++;
            }
            if (i < 0 || i >= data.size()) {
                return 0;
            }
        }
        return Double.parseDouble(applied);
    }

    public List<Data> predictData(List<?> data) {
        List<Data> result = new ArrayList<>();
        for (Object datum : data) {
            Data temp = new Data();
            BeanUtils.copyProperties(datum, temp);
            result.add(temp);
        }
        return this.predict(result);
    }

    @Override
    public void onApplicationEvent(@NotNull ContextClosedEvent event) {
        this.model.close();
    }

    @Override
    public boolean supportsAsyncExecution() {
        return ApplicationListener.super.supportsAsyncExecution();
    }

    @lombok.Data
    public static class Data {
        private String aqi;

        private String pm25;

        private String pm10;

        private String co;

        private String no2;

        private String o3;

        private String so2;
    }
}
